// #include "PCFG.h"
#include "PCFG_v2.h"
#include <chrono>
#include <fstream>
// #include "md5.h"
#include "SIMD_md5.h"
#include <iomanip>
#include <mpi.h>
using namespace std;
using namespace chrono;

// 编译指令
// mpicxx main.cpp train.cpp guessing.cpp SIMD_md5.cpp -o ../main  -fopenmp -O2

int main(int argc, char **argv)
{
  // 初始化MPI
  MPI_Init(&argc, &argv);

  int rank, size;
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
  MPI_Comm_size(MPI_COMM_WORLD, &size);
  
    double time_hash = 0;  // 用于MD5哈希的时间
    double time_guess = 0; // 哈希和猜测的总时长
    double time_train = 0; // 模型训练的总时长
    PriorityQueue q;
    double start_train = MPI_Wtime();
    q.m.train("/guessdata/Rockyou-singleLined-full.txt");
    q.m.order();
    double end_train = MPI_Wtime();
    time_train = end_train - start_train;

    q.init();
    cout << "here" << endl;
    int curr_num = 0;
    double start = MPI_Wtime();

    // cout << "Priority queue size: " << q.priority.size() << endl;
    // cout << "About to enter main loop" << endl;

    // 由于需要定期清空内存，我们在这里记录已生成的猜测总数
    int history = 0;
    // std::ofstream a("./files/results.txt");
    while (!q.priority.empty())
    {
      if(rank == 0)
      {
        q.PopNext();
        q.total_guesses = q.guesses.size();
        
        if (q.total_guesses - curr_num >= 100000)
        {
          cout << "Guesses generated: " << history + q.total_guesses << endl;
          curr_num = q.total_guesses;

          // 在此处更改实验生成的猜测上限
          int generate_n = 10000000;
          if (history + q.total_guesses > 10000000)
          {
            //发信号告诉rank1，结束程序
            int end_signal = -1;
            MPI_Send(&end_signal, 1, MPI_INT, 1, 0, MPI_COMM_WORLD); // 发送结束信号
            //当哈希结束的时候，它才会收到，这里是阻塞的，所以肯定一致
            double end = MPI_Wtime();
            time_guess = end - start;
            cout << "Guess + Hash time:" << time_guess << "seconds" << endl;
            // cout << "Hash time:" << time_hash << "seconds" << endl;
            cout << "Train time:" << time_train << "seconds" << endl;
            break;
          }
        }

        if (curr_num > 1000000){
          int start_signal = 1; // 发送哈希信号
          MPI_Send(&start_signal, 1, MPI_INT, 1, 0, MPI_COMM_WORLD); // 发送开始信号
          
          // 发送猜测字符串给另一个进程
          
          // 发送字符串个数，和总长度、发送每个字符串的长度
          int num_guesses = q.guesses.size();
          vector<int> lengths(num_guesses);
          vector<int> offsets(num_guesses); // 用于存储每个字符串的偏移量
          offsets[0] = 0; // 第一个字符串的偏移量为0
          int total_length = 0;
#pragma omp parallel for num_threads(7) shared(q, lengths) // 使用 OpenMP 并行化
          for (int i = 0; i < num_guesses; i++) {
            lengths[i] = q.guesses[i].length();
          }
          
          for (int i = 0; i < num_guesses; i++) {
            offsets[i] = total_length; // 记录每个字符串的偏移量
            total_length += lengths[i];
          }
          

          //发送总字符串
          vector<char> all_guesses(total_length + 1);
          all_guesses[total_length] = '\0';
#pragma omp parallel for num_threads(7) shared(all_guesses,q, offsets) // 使用 OpenMP 并行化
          for (int i = 0; i < num_guesses; i++)
          {
            memcpy(&all_guesses[offsets[i]], q.guesses[i].data(), lengths[i]);
          }

          // 发给进程1，一起发
          MPI_Send(&num_guesses, 1, MPI_INT, 1, 0, MPI_COMM_WORLD);            // 发送字符串个数
          MPI_Send(&total_length, 1, MPI_INT, 1, 0, MPI_COMM_WORLD);           // 发送总长度
          MPI_Send(lengths.data(), num_guesses, MPI_INT, 1, 0, MPI_COMM_WORLD);       // 发送每个字符串的长度
          MPI_Send(all_guesses.data(), total_length, MPI_CHAR, 1, 0, MPI_COMM_WORLD); // 发送所有字符串
          
          
          //恢复数据，记录已经生成的口令总数
          history += curr_num;
          curr_num = 0;
          q.guesses.clear();
        }
      }
      
      
      else if(rank ==1){
        
      // 接收开始或结束信息
      int signal;
      MPI_Recv(&signal, 1, MPI_INT, 0, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
      if (signal == -1)
        break;
        
      // 接收字符串个数和总长度
      int num_guesses, total_length;
      MPI_Recv(&num_guesses, 1, MPI_INT, 0, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
      MPI_Recv(&total_length, 1, MPI_INT, 0, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
      // 接收每个字符串的长度
      vector<int> lengths(num_guesses);
      MPI_Recv(lengths.data(), num_guesses, MPI_INT, 0, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
      // 接收所有字符串
      vector<char> all_guesses(total_length+1);
      MPI_Recv(all_guesses.data(), total_length, MPI_CHAR, 0, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);

      all_guesses[total_length] = '\0'; // 确保字符串以 null 结尾
      
      // 计算偏移量
      vector<int> offsets(num_guesses); // 用于存储每个字符串的偏移量
      offsets[0] = 0; // 第一个字符串的偏移量为0
      for(int i = 0; i < num_guesses-1; i++) {
        offsets[i + 1] = offsets[i] + lengths[i]; // 计算每个字符串的偏移量
      }

  // 使用 OpenMP 并行化外层循环
#pragma omp parallel for num_threads(7) shared(offsets,all_guesses)
      for (int i = 0; i < num_guesses; i += 4)
      {
        alignas(16) bit32 state[4][4];
        string pw[4];
        for(int j = 0; j < 4 && i + j < num_guesses; j++) {
          pw[j] = string(&all_guesses[offsets[i + j]], lengths[i + j]); // 使用偏移量获取每个字符串

        }
        SIMD_MD5Hash(pw, state); // 假设 SIMD_MD5Hash 是线程安全的
      }
      

      }
      else{//我们只用两个线程
        break;
      }
  }
  
  // 完成MPI通信
  MPI_Finalize();
  return 0;
}